import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.animation import FuncAnimation
import heapq

# 参数设置（来自论文）
V0 = 1.2  # 初始速度 (m/s)
a = 0.71
b = 0.51
c_down = 1  # 下楼时的体力消耗系数
c_up = 2.7  # 上楼时的体力消耗系数

# 三维网格参数
NUM_FLOORS = 10  # 楼层数
FLOOR_HEIGHT = 3  # 每层楼高度（米）
STEPS_PER_FLOOR = 20  # 每层楼网格数
GRID_SIZE = 0.5  # 每个网格的边长（米）

# 初始化人群参数
NUM_PEOPLE = 50  # 总人数
START_FLOORS = np.random.randint(2, NUM_FLOORS, size=NUM_PEOPLE)  # 避免初始在底层

class BuildingGrid3D:
    def __init__(self, num_floors, floor_height, steps_per_floor):
        self.num_floors = num_floors
        self.floor_height = floor_height
        self.steps_per_floor = steps_per_floor
        self.grid_shape = (num_floors, steps_per_floor, steps_per_floor)
        self.exit_location = (0, steps_per_floor//2, steps_per_floor//2)  # 出口在底层中心
        
        # 初始化障碍物网格（走廊可通行，房间为障碍物）
        self.obstacle_grid = np.ones(self.grid_shape, dtype=bool)  # 默认为障碍物
        # 设置走廊区域（中心区域）
        corridor_width = 4
        for floor in range(num_floors):
            self.obstacle_grid[floor,
                              steps_per_floor//2 - corridor_width:steps_per_floor//2 + corridor_width,
                              steps_per_floor//2 - corridor_width:steps_per_floor//2 + corridor_width] = False
            
            # 设置楼梯位置（中心点）
            self.obstacle_grid[floor, steps_per_floor//2, steps_per_floor//2] = False
            # 楼梯连接下层
            if floor < num_floors - 1:
                self.obstacle_grid[floor+1, steps_per_floor//2, steps_per_floor//2] = False
        
        self.static_field = None  # 静态场将在此存储

    def get_neighbors_36(self, pos):
        """获取36邻域的邻居（考虑障碍物）"""
        x, y, z = pos
        neighbors = []
        for dx in [-1, 0, 1]:
            for dy in [-1, 0, 1]:
                for dz in [-1, 0, 1]:
                    if dx == 0 and dy == 0 and dz == 0:
                        continue  # 排除自身
                    nx, ny, nz = x + dx, y + dy, z + dz
                    # 检查是否在网格范围内
                    if (0 <= nx < self.num_floors and 
                        0 <= ny < self.steps_per_floor and 
                        0 <= nz < self.steps_per_floor):
                        # 检查是否是障碍物
                        if not self.obstacle_grid[nx, ny, nz]:
                            neighbors.append((nx, ny, nz))
        return neighbors

    def is_valid_move(self, current_pos, next_pos):
        """人群移动时只能向下或同层，且楼梯位置允许下楼"""
        current_floor = current_pos[0]
        next_floor = next_pos[0]
        # 检查是否在楼梯位置（中心点）
        is_stair = (current_pos[1] == self.steps_per_floor//2 and 
                    current_pos[2] == self.steps_per_floor//2)
        
        if next_floor < current_floor:  # 下楼
            return is_stair
        elif next_floor == current_floor:  # 同层移动
            return True
        else:  # 不允许上楼
            return False

def generate_static_field(grid, exit_pos):
    """生成静态场：计算每个位置到出口的最短距离（考虑障碍物）"""
    inf = float('inf')
    dist = np.full(grid.grid_shape, inf)
    dist[exit_pos] = 0
    visited = np.zeros(grid.grid_shape, dtype=bool)
    
    heap = []
    heapq.heappush(heap, (0, exit_pos))
    
    while heap:
        current_dist, current_pos = heapq.heappop(heap)
        if visited[current_pos]:
            continue
        visited[current_pos] = True
        
        for neighbor in grid.get_neighbors_36(current_pos):
            if not grid.obstacle_grid[neighbor]:
                new_dist = current_dist + 1
                if new_dist < dist[neighbor]:
                    dist[neighbor] = new_dist
                    heapq.heappush(heap, (new_dist, neighbor))
    
    # 验证顶层中心位置的静态场值
    test_pos = (grid.num_floors-1, grid.steps_per_floor//2, grid.steps_per_floor//2)
    print(f"顶层中心位置静态场值: {dist[test_pos]}")
    return dist

class Person:
    def __init__(self, pos, grid):
        self.pos = pos
        self.grid = grid
        self.v0 = V0
        self.sum_d = 0
        self.sum_t = 0
        self.time_step = 0.5  # 缩短时间步长提升流畅度
        self.countdown = 0

    def update_velocity(self):
        current_floor = self.pos[0]
        c = c_down if current_floor > 0 else c_up
        exponent = -0.01 * c
        velocity = self.v0 * (1 - a * self.sum_d - b * self.sum_t) * np.exp(exponent)
        velocity = max(velocity, 0.1)
        self.time_step = 1.0 / velocity

    def move(self):
        current_dist = self.grid.static_field[self.pos]
        if current_dist == float('inf'):
            return
        
        best_next = None
        min_dist = float('inf')
        for neighbor in self.grid.get_neighbors_36(self.pos):
            if self.grid.is_valid_move(self.pos, neighbor):
                if self.grid.obstacle_grid[neighbor]:
                    continue  # 障碍物不可通行
                neighbor_dist = self.grid.static_field[neighbor]
                if neighbor_dist < min_dist:
                    min_dist = neighbor_dist
                    best_next = neighbor
        
        if best_next:
            self.pos = best_next
            self.sum_d += 1
            self.sum_t += self.time_step

def main():
    grid = BuildingGrid3D(NUM_FLOORS, FLOOR_HEIGHT, STEPS_PER_FLOOR)
    static_field = generate_static_field(grid, grid.exit_location)
    grid.static_field = static_field
    
    # 验证静态场
    test_pos = (NUM_FLOORS-1, STEPS_PER_FLOOR//2, STEPS_PER_FLOOR//2)
    assert static_field[test_pos] != float('inf'), "顶层中心位置应可达出口"
    
    # 初始化人群（在走廊区域随机生成）
    people = []
    for floor in START_FLOORS:
        # 随机生成走廊区域坐标
        y = np.random.randint(grid.steps_per_floor//2 - 2, grid.steps_per_floor//2 + 3)
        z = np.random.randint(grid.steps_per_floor//2 - 2, grid.steps_per_floor//2 + 3)
        pos = (floor, y, z)
        person = Person(pos, grid)
        people.append(person)

    # 可视化设置
    fig = plt.figure(figsize=(12, 8))
    ax = fig.add_subplot(111, projection='3d')
    ax.set_xlabel('Floor')
    ax.set_ylabel('X')
    ax.set_zlabel('Y')
    ax.set_xlim(0, NUM_FLOORS)
    ax.set_ylim(0, grid.steps_per_floor)
    ax.set_zlim(0, grid.steps_per_floor)

    # 动态更新函数（每帧执行5次移动）
    def update(frame):
        for _ in range(5):  # 每帧执行5次移动
            for person in people:
                person.countdown += 1
                if person.countdown >= person.time_step:
                    person.update_velocity()
                    person.move()
                    person.countdown = 0

        ax.cla()
        ax.set_xlabel('Floor')
        ax.set_ylabel('X')
        ax.set_zlabel('Y')
        ax.set_xlim(0, NUM_FLOORS)
        ax.set_ylim(0, grid.steps_per_floor)
        ax.set_zlim(0, grid.steps_per_floor)

        # 绘制人群位置
        xs = [p.pos[0] for p in people]
        ys = [p.pos[1] for p in people]
        zs = [p.pos[2] for p in people]
        ax.scatter(xs, ys, zs, c='red', s=50, label='People')

        # 绘制出口
        ax.scatter(*grid.exit_location, c='green', s=100, label='Exit')

        # 绘制楼梯（中心点）
        for floor in range(grid.num_floors):
            ax.scatter(floor, grid.steps_per_floor//2, grid.steps_per_floor//2, c='gray', s=10)

    # 创建动画（更高帧率）
    ani = FuncAnimation(fig, update, frames=200, interval=50)
    plt.show()

if __name__ == "__main__":
    main()